begin

    using BenchmarkTools

    using CSV

    using LaTeXStrings

    using NearestNeighbors

    using Plots

    using PlutoUI

    using Random

end

br = html"<br>";

bigbr = html"<br><br><br><br><br>";

html"<button onclick=present()>Present</button>"

@enum InfectionStatus S I R

R

"""

Initialize arrays of walkers and their statuses

​

    N - number of agents per cluster

    centers - Center points of clusters, Tuple of tuples

    L - Box size for simulation

"""

function initialize(N, L, rng)

    positions = rand(rng, 1-L:L-1, (2, N))  # Spread agents out across board

    positions = positions .+ (rand(rng, 2, N) .- 0.5)  # Give each agent a unique pos

    states = [S for i = 1:N]

    states[1] = I

    return positions, states

end;

"""

Walk a single point by a random amount.

"""

function walk(x, y, L, rng)

    i = randn(rng)

    j = randn(rng)

    # move walkers by random amount

    x += i

    y += j

    # if the posision exceeds the boundary L, then move back

    if L ≤ abs(x)  || L ≤ abs(y)

        x -= i

        y -= j

    end

    return x, y

end ;

"""

Walk method that accepts a 1x2 array or a 2-tuple as input.

"""

walk(array::Union{Array, Tuple}, L, rng) = walk(array[1], array[2], L, rng);

"""

Get the trajectory of a random walker over N time steps.

"""

function trajectory(N, L=200, rng=MersenneTwister(convert(Int, 1e9*time())))

    # Pre-allocate array for efficiency

    path = Array{Tuple{Float64, Float64}, 1}(undef, N)

    # Start at the origin

    path[1] = (0.0, 0.0)

    # Do the walk

    for i = 2:N

        path[i] = walk(path[i-1], L, rng)

    end

    return path

end;

# Make an array of 10 walkers

walkers = [trajectory(20_000) for i in 1:10]

"""

Convenience function combining `BallTee` creation and `inrange` functions.

"""

encounters(contacts, individual, radius) = inrange(BallTree(contacts), individual, radius);

function step!(positions, states, N, L, pᵢ, pᵣ, radius, rng::MersenneTwister)

    # Pick an agent at random

    i = rand(rng, 1:N)

    # Give infected agents a chance to recover before infecting others

    if states[i] == I

        if rand(rng) ≤ pᵣ

            states[i] = R

        end

    end

    # Move the agent to a new location

    xy = walk(positions[:, i], L, rng)

    positions[1, i]  = xy[1]

    positions[2, i]  = xy[2]

    # If the agent is still infected

    if states[i] == I

        # Determine close contacts

        contacts = encounters(positions, positions[:, i], radius)

        # Remove self

        contacts = contacts[(contacts .≠ i)]

        # For each contact, 

        for contact in contacts

            # Give a chance to infect if susceptible

            if states[contact] == S

                if rand(rng) ≤ pᵢ

                    states[contact] = I

                end

            end

        end

    end

    return positions, states

end;

function sweep!(positions, states, N, L, pᵢ, pᵣ, radius, rng::MersenneTwister)

    for n = 1:N

        step!(positions, states, N, L, pᵢ, pᵣ, radius, rng)

    end

    return positions, states

end;

function simulation(Sw, N, L, pᵢ, pᵣ, radius, rng=MersenneTwister(convert(Int, 1e9*time())))

    system = Dict(

          "S" => Sw,  # Number of sweeps ("days")

          "N" => N,  # Number of agents

          "L" => L,  # Length of bounding box for agents

          "pᵢ" => pᵢ,  # Infection probability

          "pᵣ" => pᵣ,  # Recovery probability

          "positions" => [],  # Positions per sweep

          "states" => [],  # States per sweep

          "susceptible" => Int[],  # Number of susceptible agents at step n

          "infected" => Int[],  # Number of infected agents at step n

          "recovered" => Int[])  # Number of recovered agents at step n

    positions, states = initialize(N, L, rng)

    # Large sim threshold for saving movement history

    thresh = 1e5

    if Sw*N < thresh

        push!(system["positions"], deepcopy(positions))

    end

    push!(system["states"], deepcopy(states))

    push!(system["susceptible"], sum(states .== S))

    push!(system["infected"], sum(states .== I))

    push!(system["recovered"], sum(states .== R))

    for n=2:Sw

        sweep!(positions, states, N, L, pᵢ, pᵣ, radius, rng)

        if Sw*N < thresh

            push!(system["positions"], deepcopy(positions))

        end

        push!(system["states"], deepcopy(states))

        push!(system["susceptible"], sum(states .== S))

        push!(system["infected"], sum(states .== I))

        push!(system["recovered"], sum(states .== R))

    end

    return system

end;

med_sim = simulation(100, 800, 20, 0.3, 0.05, 1);

# "Big" Simulation parameters

begin

    days = 500

    N = 1500

    boxsize = 30 

    prob_i = 0.1

    prob_r = 0.02

    cluster_rad = 2.5

end;